# ============================================================
# ssq_transformer_full.py  (双色球版本 2025-08-29)
# ============================================================
import os, json, math, random
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# ------------------------------------------------------------
# 0. 超参
# ------------------------------------------------------------
SEQ_LEN     = 100          # 用最近 100 期预测下一期
RED_BALLS   = 33          # 红球 1-33
BLUE_BALLS  = 16          # 蓝球 1-16
BATCH_SIZE  = 32
EPOCHS      = 50
LR          = 1e-3
DEVICE      = 'cuda' if torch.cuda.is_available() else 'cpu'
CSV_DIR     = 'data'   # 数据目录
# 获取所有CSV文件，确保目录存在且文件可读
CSV_FILES = []
if os.path.exists(CSV_DIR) and os.path.isdir(CSV_DIR):
    CSV_FILES = sorted([
        os.path.join(CSV_DIR, f) 
        for f in os.listdir(CSV_DIR) 
        if f.endswith('.csv') and os.path.isfile(os.path.join(CSV_DIR, f))
    ])
else:
    print(f"错误：目录 {CSV_DIR} 不存在或不是目录")

# ------------------------------------------------------------
# 1. 读取 & 清洗
# ------------------------------------------------------------
def load_ssq(csv_path: str):
    """
    解析形如：
    期号,开奖日期,红球,蓝球,一等奖注数,...
    2025099,2025-08-28(四),09 11 15 17 22 26,14,...
    """
    df = pd.read_csv(csv_path)
    # 日期统一为 datetime
    df['date'] = pd.to_datetime(df['开奖日期'].str.split('(').str[0])
    # 拆红球
    red_df = (
        df['红球']
        .str.split(expand=True)
        .astype(int)
        .rename(columns={i: f'r{i+1}' for i in range(6)})
    )
    # 蓝球
    blue_df = df[['蓝球']].astype(int).rename(columns={'蓝球': 'b'})
    return pd.concat([df[['期号', 'date']], red_df, blue_df], axis=1)

# 合并所有CSV文件的数据
dfs = []
for csv_file in CSV_FILES:
    try:
        dfs.append(load_ssq(csv_file))
    except Exception as e:
        print(f"加载文件 {csv_file} 时出错: {str(e)}")
df = pd.concat(dfs, ignore_index=True).sort_values('date')
print('总样本数:', len(df))
print(df.head())

# ------------------------------------------------------------
# 2. one-hot 编码
#   X_red : (N, 33)   红球
#   X_blue: (N, 16)   蓝球
# ------------------------------------------------------------
X_red  = np.zeros((len(df), RED_BALLS),  dtype=np.float32)
X_blue = np.zeros((len(df), BLUE_BALLS), dtype=np.float32)

for i, row in df.iterrows():
    reds  = row[['r1','r2','r3','r4','r5','r6']].values.astype(int)
    blue  = int(row['b'])
    X_red[i,  reds-1] = 1
    X_blue[i, blue-1] = 1

# ------------------------------------------------------------
# 3. 数据集：同时返回红球和蓝球标签
# ------------------------------------------------------------
class SsqDataset(Dataset):
    def __init__(self, X_red, X_blue, seq_len=SEQ_LEN):
        self.Xr, self.Xb, self.seq_len = X_red, X_blue, seq_len
    def __len__(self):
        return len(self.Xr) - self.seq_len
    def __getitem__(self, idx):
        xr = torch.tensor(self.Xr[idx:idx+self.seq_len])
        xb = torch.tensor(self.Xb[idx:idx+self.seq_len])
        yr = torch.tensor(self.Xr[idx+self.seq_len])
        yb = torch.tensor(self.Xb[idx+self.seq_len])
        return xr, xb, yr, yb

ds     = SsqDataset(X_red, X_blue)
# 每次迭代时从数据集中加载的样本数量 batch_size
# 训练阶段True 每个epoch随机打乱数据顺序 False 验证/测试阶段 保持数据原始顺序
loader = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True)

# ------------------------------------------------------------
# 4. 双色球 Transformer
#   - 输入: (B, T, 49)  33+16 拼接
#   - 输出: 红球 logits (B, 33) + 蓝球 logits (B, 16)
# ------------------------------------------------------------
class SsqTransformer(nn.Module):
    def __init__(self, d_model=49, nhead=7, nlayers=4):
        super().__init__()
        self.pos = nn.Parameter(torch.randn(SEQ_LEN, d_model))
        encoder = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=256, batch_first=True
        )
        self.tf  = nn.TransformerEncoder(encoder, num_layers=nlayers)
        self.fc_r = nn.Linear(d_model, RED_BALLS)
        self.fc_b = nn.Linear(d_model, BLUE_BALLS)

    def forward(self, xr, xb):
        x = torch.cat([xr, xb], dim=-1)          # (B, T, 49)
        x = x + self.pos
        x = self.tf(x)[:, -1, :]                 # (B, 49)
        return torch.sigmoid(self.fc_r(x)), torch.sigmoid(self.fc_b(x))

model = SsqTransformer().to(DEVICE)
opt   = torch.optim.AdamW(model.parameters(), lr=LR)
loss_fn = nn.BCELoss()

# ------------------------------------------------------------
# 5. 训练
# ------------------------------------------------------------
model.train()
for epoch in range(1, EPOCHS+1):
    total = 0
    for xr, xb, yr, yb in loader:
        xr, xb, yr, yb = xr.to(DEVICE), xb.to(DEVICE), yr.to(DEVICE), yb.to(DEVICE)
        pr, pb = model(xr, xb)
        loss = loss_fn(pr, yr) + loss_fn(pb, yb)
        opt.zero_grad(); loss.backward(); opt.step()
        total += loss.item()
    if epoch % 10 == 0 or epoch == 1:
        print(f'Epoch {epoch:02d} | loss {total/len(loader):.4f}')

# ------------------------------------------------------------
# 6. 预测下一期
# ------------------------------------------------------------
# model.eval()
# with torch.no_grad():
#     xr_last = torch.tensor(X_red[-SEQ_LEN:]).unsqueeze(0).to(DEVICE)
#     xb_last = torch.tensor(X_blue[-SEQ_LEN:]).unsqueeze(0).to(DEVICE)
#     pr, pb  = model(xr_last, xb_last)

#     # 红球：取概率最高的 6 个
#     top_red   = torch.topk(pr[0], k=6).indices.cpu().numpy() + 1
#     # 蓝球：取概率最高的 1 个
#     top_blue  = torch.topk(pb[0], k=1).indices.cpu().numpy() + 1

# print('\n=== 下一期预测 ===')
# print('红球:', sorted(top_red))
# print('蓝球:', int(top_blue))

# ------------------------------------------------------------
# 7. 回测最近 100 期
# ------------------------------------------------------------
def evaluate(model, Xr, Xb, last_n=100):
    hits_red, hits_blue = 0, 0
    tot_red,  tot_blue  = last_n*6, last_n*1
    model.eval()
    with torch.no_grad():
        for i in range(len(Xr)-SEQ_LEN-last_n, len(Xr)-SEQ_LEN):
            xr = torch.tensor(Xr[i:i+SEQ_LEN]).unsqueeze(0).to(DEVICE)
            xb = torch.tensor(Xb[i:i+SEQ_LEN]).unsqueeze(0).to(DEVICE)
            pr, pb = model(xr, xb)
            pred_red  = torch.topk(pr[0], k=6).indices.cpu().numpy()
            pred_blue = torch.topk(pb[0], k=1).indices.cpu().numpy()
            true_red  = np.where(Xr[i+SEQ_LEN] == 1)[0]
            true_blue = np.where(Xb[i+SEQ_LEN] == 1)[0]
            hits_red  += len(set(pred_red)  & set(true_red))
            hits_blue += len(set(pred_blue) & set(true_blue))
    print(f'回测最近 {last_n} 期')
    print(f'  红球命中 {hits_red}/{tot_red} ≈ {hits_red/tot_red:.3f}')
    print(f'  蓝球命中 {hits_blue}/{tot_blue} ≈ {hits_blue/tot_blue:.3f}')
    return hits_red, hits_blue, tot_red, tot_blue

# hits_red, hits_blue, tot_red, tot_blue = evaluate(model, X_red, X_blue)
# print("red_accuracy is:{0:.3f}, blue_accuracy is:{1:.3f}".format(hits_red/tot_red, hits_blue/tot_blue))

# ------------------------------------------------------------
# 8. 循环运行并保存结果
# ------------------------------------------------------------
all_qualified_predictions = []
for run in range(1):
    # 重新训练模型
    model = SsqTransformer().to(DEVICE)
    opt = torch.optim.AdamW(model.parameters(), lr=LR)
    
    model.train()
    for epoch in range(1, EPOCHS+1):
        total = 0
        for xr, xb, yr, yb in loader:
            xr, xb, yr, yb = xr.to(DEVICE), xb.to(DEVICE), yr.to(DEVICE), yb.to(DEVICE)
            pr, pb = model(xr, xb)
            loss = loss_fn(pr, yr) + loss_fn(pb, yb)
            opt.zero_grad(); loss.backward(); opt.step()
            total += loss.item()
    
    # 预测下一期
    model.eval()
    with torch.no_grad():
        xr_last = torch.tensor(X_red[-SEQ_LEN:]).unsqueeze(0).to(DEVICE)
        xb_last = torch.tensor(X_blue[-SEQ_LEN:]).unsqueeze(0).to(DEVICE)
        pr, pb = model(xr_last, xb_last)
        top_red = torch.topk(pr[0], k=6).indices.cpu().numpy() + 1
        top_blue = torch.topk(pb[0], k=1).indices.cpu().numpy() + 1
    
    # 评估
    hits_red, hits_blue, tot_red, tot_blue = evaluate(model, X_red, X_blue)
    
    # 检查条件并保存结果
    if hits_blue/tot_blue == 1.0 and hits_red/tot_red > 0.5:
        prediction_data = {
            'run_id': run,
            'time': pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S'),
            'red_balls': [int(x) for x in sorted(top_red)],
            'blue_ball': int(top_blue),
            'red_accuracy': hits_red/tot_red,
            'blue_accuracy': hits_blue/tot_blue
        }
        all_qualified_predictions.append(prediction_data)
        print(f'第{run}次运行结果符合条件')

# 保存所有符合条件的结果
if all_qualified_predictions:
    json.dump(all_qualified_predictions, open('all_qualified_predictions.json', 'w'), indent=2)
    print(f'共找到{len(all_qualified_predictions)}组符合条件的预测结果，已保存到all_qualified_predictions.json')
else:
    print('未找到符合条件的预测结果')

# 保存最终模型和概率数据
torch.save(model.state_dict(), 'ssq_transformer.pt')
prob_data = {
    'red_prob': torch.softmax(pr[0], -1).cpu().tolist(),
    'blue_prob': torch.softmax(pb[0], -1).cpu().tolist()
}
json.dump(prob_data, open('ssq_next_probs.json', 'w'), indent=2)

print('模型与概率已保存')